// SPDX-License-Identifier: MIT

#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/ec.h>
#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/provider.h>
#include <oqs/oqs.h>
#include <pthread.h>
#if defined(OQS_USE_RASPBERRY_PI)
#define _OQS_RASPBERRY_PI
#endif
#if defined(OQS_SPEED_USE_ARM_PMU)
#define SPEED_USE_ARM_PMU
#endif
#include "ds_benchmark.h"
#include "system_info.c"
#include "test_common.h"

struct SigPair {
    unsigned char *sigma1;
    size_t sigma1_len;
    unsigned char *sigma2;
    size_t sigma2_len;
} signature;

typedef struct {
    const OQS_SIG *sig;
    const uint8_t *message;
    size_t message_len;
    const uint8_t *new_message;
    size_t new_message_len;
    const uint8_t *public_key;
  EVP_PKEY *classpkey;
    int *result;
} VerifyArgs;

const char *kHybridSignatureAlgorithms[] = {
    "p256_dilithium2", "rsa3072_dilithium2", "p384_dilithium3", "p521_dilithium5",
    "p256_ML-DSA-44", "rsa3072_ML-DSA-44", "p384_ML-DSA-65", "p521_ML-DSA-87",
    "p256_Falcon-512", "rsa3072_Falcon-512", "p256_Falcon-padded-512", "rsa3072_Falcon-padded-512",
    "p521_Falcon-1024", "p521_Falcon-padded-1024", "p256_SPHINCS+-SHA2-128f-simple", "rsa3072_SPHINCS+-SHA2-128f-simple",
    "p256_SPHINCS+-SHA2-128s-simple", "rsa3072_SPHINCS+-SHA2-128s-simple", "p384_SPHINCS+-SHA2-192f-simple",
    "p256_SPHINCS+-SHAKE-128f-simple", "rsa3072_SPHINCS+-SHAKE-128f-simple", NULL
};
/*
const char *kHybridSignatureAlgorithms[] = {
    "p256_Falcon-padded-512", NULL
};*/
static OSSL_LIB_CTX *init_openssl(void) {
    OSSL_LIB_CTX *ctx;

    if (!(ctx = OSSL_LIB_CTX_new())) {
        fprintf(stderr, cRED "failed to initialize a new `OSSL_LIB_CTX` " cNORM "\n");
    }
    return ctx;
}

char *get_oqs_fromtls(const char *tlsname){
    char *oqsname;
    char *del = "_";
    char *position = strchr(tlsname, *del);
    
    if(position == NULL)    return NULL;
    
  oqsname = position+1;    
    return oqsname;
}

char *get_cls_fromtls(const char *tlsname) {
     char *clsname;
     char *del = "_";
     char *position = strchr(tlsname, *del);

     size_t len = position - tlsname;
  clsname = (char *)malloc(len + 1);  // +1 为了 '\0'
    if (clsname == NULL)    return NULL;
    
  strncpy(clsname, tlsname, len);
  clsname[len] = '\0';
    return clsname;
}

static EVP_PKEY *generate_classic_evp(const char *classic_names){
         
     EVP_PKEY *pkey = NULL;
            if(!strncmp(classic_names,"rsa", 3)){
        pkey = EVP_RSA_gen(3072);
                    return pkey;
     }
           if(!strncmp(classic_names, "p256",4)){
         pkey =  EVP_EC_gen("P-256");
                       return pkey;
     }
            if(!strncmp(classic_names, "p384",4)){
         pkey =  EVP_EC_gen("P-384");
                       return pkey;
     }        
            if(!strncmp(classic_names, "p521",4)){
         pkey =  EVP_EC_gen("P-521");
                        return pkey;
     }          
    EVP_PKEY_free(pkey);
           return NULL;
}


static OQS_STATUS verify_hybrid_signature(const OQS_SIG *sig, EVP_PKEY *classpkey, const uint8_t *message, size_t message_len,const uint8_t *new_message, size_t new_message_len, const uint8_t *public_key){
   EVP_MD_CTX *mdctx = EVP_MD_CTX_new(); 
          if(!mdctx || !EVP_DigestVerifyInit(mdctx, NULL, EVP_sha256(), NULL, classpkey) || !EVP_DigestVerifyUpdate(mdctx, message, message_len) || !EVP_DigestVerifyFinal(mdctx, signature.sigma1, signature.sigma1_len)){
          fprintf(stderr, "Classic pkey verify failed.\n");
          EVP_MD_CTX_free(mdctx);
                           return OQS_ERROR;
    }
          // int ret;
    OQS_SIG_verify(sig, new_message, new_message_len, signature.sigma2, signature.sigma2_len, public_key);
    //printf("%d ", ret);
    EVP_MD_CTX_free(mdctx);
          return OQS_SUCCESS;
}



static OQS_STATUS generate_hybrid_signature(const OQS_SIG *sig, EVP_PKEY *classpkey, const uint8_t *message, size_t message_len, const uint8_t *secret_key){
    EVP_MD_CTX *mdctx = EVP_MD_CTX_new(); 
          if(!mdctx ||!EVP_DigestSignInit(mdctx, NULL, EVP_sha256(), NULL, classpkey) ||!EVP_DigestSignUpdate(mdctx, message, message_len) ||!EVP_DigestSignFinal(mdctx, NULL, &signature.sigma1_len)|| (signature.sigma1 = OPENSSL_malloc(signature.sigma1_len)) == NULL ||!EVP_DigestSignFinal(mdctx, signature.sigma1, &signature.sigma1_len)){
          fprintf(stderr, "Classic pkey signed failed.\n");
          EVP_MD_CTX_free(mdctx);
                           return OQS_ERROR;
    }
           uint8_t *new_message = NULL;
           size_t new_message_len = signature.sigma1_len + message_len;
    new_message = OPENSSL_malloc(new_message_len + 1);
    memcpy(new_message, signature.sigma1, signature.sigma1_len);
    memcpy(new_message + signature.sigma1_len, message, message_len);
    new_message[new_message_len] = '\0';
    OQS_SIG_sign(sig, signature.sigma2, &signature.sigma2_len, new_message, new_message_len, secret_key);
 //   fprintf(stderr, "Signature length is %d.\n", signature.sigma1_len);
    EVP_MD_CTX_free(mdctx); 
    OQS_MEM_insecure_free(new_message);	   
          return OQS_SUCCESS;          
}

static OQS_STATUS test_algorithm(OSSL_LIB_CTX *libctx,const char *signames, uint64_t duration, bool printInfo, bool doFullCycle){

	OQS_SIG *sig = NULL;
	uint8_t *public_key = NULL;
	uint8_t *secret_key = NULL;
	uint8_t *message = NULL;
	size_t message_len = 5000;
	OQS_STATUS ret = OQS_ERROR;	

   signature.sigma1 = NULL;
   signature.sigma1_len = 0;
   signature.sigma2 = NULL;
   signature.sigma2_len = 0;
   
        char *oqsname, *clasname;	
   oqsname = get_oqs_fromtls((char *)signames);
   clasname = get_cls_fromtls((char *)signames);
//fprintf(stderr, "%s, %s.\n", clasname, oqsname);
   
   EVP_PKEY *class_pkey = NULL;
   class_pkey  = generate_classic_evp(clasname);

	sig = OQS_SIG_new(oqsname);
	if (sig == NULL) {
		goto err;
	}

	public_key = malloc(sig->length_public_key);
	secret_key = malloc(sig->length_secret_key);
	message = malloc(message_len);
	signature.sigma2_len = sig->length_signature;
	signature.sigma2 = malloc(sig->length_signature);
	
	if ((public_key == NULL) || (secret_key == NULL) || (message == NULL) || (signature.sigma2 == NULL)) {
		fprintf(stderr, "ERROR: malloc failed\n");
		goto err;
	}	
	
	OQS_randombytes(message, message_len);
   //OQS_SIG_keypair(sig, public_key, secret_key);
   
	printf("|%-36s | %10s | %14s | %15s | %10s | %25s | %10s|\n", signames, "", "", "", "", "", "");
       if (!doFullCycle) {
      TIME_OPERATION_SECONDS((generate_classic_evp(clasname)&&OQS_SIG_keypair(sig, public_key, secret_key)), "keypair", duration)
		TIME_OPERATION_SECONDS(generate_hybrid_signature(sig, class_pkey, message, message_len, secret_key), "sign", duration)
		uint8_t *new_message = NULL;
               size_t new_message_len = signature.sigma1_len + message_len;
      new_message = OPENSSL_malloc(new_message_len + 1);
      memcpy(new_message, signature.sigma1, signature.sigma1_len);
      memcpy(new_message + signature.sigma1_len, message, message_len);
      new_message[new_message_len] = '\0';
      
//   TIME_OPERATION_SECONDS(verify_hybrid_signature(sig, class_pkey, message, message_len, new_message, new_message_len, secret_key), "verify", duration)
//	OQS_SIG_sign(sig, signature, &signature_len, message, message_len, secret_key), "sign", duration)
 		TIME_OPERATION_SECONDS(OQS_SIG_verify(sig, new_message, new_message_len, signature.sigma2, signature.sigma2_len, public_key), "verify", duration)
	}

	 if (printInfo) {
		 printf("|signature bytes:%-21s| %10zu | %14s | %15s | %10s | %25s | %10s|\n", "", signature.sigma1_len + signature.sigma2_len,"", "", "", "", "");
	}	 
	
   OQS_MEM_secure_free(secret_key, sig->length_secret_key);
	OQS_MEM_insecure_free(public_key);
	OQS_MEM_insecure_free(message);	
	OQS_SIG_free(sig);
	return OQS_SUCCESS;
err:
	if (sig != NULL) {
		OQS_MEM_secure_free(secret_key, sig->length_secret_key);
	}
	OQS_MEM_insecure_free(public_key);
	OQS_MEM_insecure_free(message);		
	OQS_SIG_free(sig);
	return OQS_ERROR;
}


int main(int argc, char **argv) {

	int ret = EXIT_SUCCESS;
	OSSL_LIB_CTX *libctx;
   OSSL_PROVIDER *oqs_provider;
	const char *modulename;
        const char *configfile;
	bool printUsage = false;
	uint64_t duration = 5;
	bool printSigInfo = true;
	bool doFullCycle = false;

       if (!(libctx = init_openssl())) {
                   goto end;
   }	
	T(argc == 3);
   modulename = argv[1];
   configfile = argv[2];
	
        if(!(oqs_provider = OSSL_PROVIDER_load(libctx, modulename))){
        fputs(cRED "  `oqs_provider` is NULL " cNORM "\n", stderr);
                      goto end;    
    }	
	
   //sigalgs = OSSL_PROVIDER_query_operation(oqs_provider, OSSL_OP_SIGNATURE, &query_nocache);		
		
   print_system_info();

	printf("Speed test\n");
	printf("==========\n");

	PRINT_TIMER_HEADER	
       for (int i = 0 ; kHybridSignatureAlgorithms[i] != NULL; i++) {
                 if (test_algorithm(libctx, kHybridSignatureAlgorithms[i], duration, printSigInfo, doFullCycle)) {
            fprintf(stderr, cRED " Failed for %s " cNORM "\n", kHybridSignatureAlgorithms[i]);
       } else {
            //fprintf(stderr, cGREEN "%s succeeded" cNORM "\n", kHybridSignatureAlgorithms[i]);       
       }
   } 	

	 PRINT_TIMER_FOOTER
end:
    OSSL_PROVIDER_unload(oqs_provider);   
    OSSL_LIB_CTX_free(libctx);
	return ret;
}
